import os
import yaml
import json

##==============================================================================
# 从用户配置中读取数据集配置，并转换为mmdetection的格式
def load_dataset_cfgs( user_cfg_path ):
    with open(user_cfg_path) as f:
        user_cfg = yaml.load(f, Loader=yaml.Loader)

    if user_cfg['dataset_type'] == 'VOC':
        return load_voc_dataset( user_cfg )
    elif user_cfg['dataset_type'] == 'COCO':
        return load_coco_dataset( user_cfg )
    else:
        return None

def load_voc_dataset( user_cfg ):
    # 读取数据集的类别信息
    classes = user_cfg['classes']
    num_classes = len(classes)

    # 生成mmdetection对应的字典格式
    # 设置训练集相关的配置
    cfg_dict = {}
    cfg_dict['TrainDataset'] = {}
    cfg_dict['TrainDataset']['name'] = 'VOCDataSet'
    cfg_dict['TrainDataset']['dataset_dir'] = user_cfg['TrainDataset']['dataset_dir']
    cfg_dict['TrainDataset']['anno_path'] = user_cfg['TrainDataset']['anno_path']
    cfg_dict['TrainDataset']['label_list'] = user_cfg['TrainDataset']['label_list']
    cfg_dict['TrainDataset']['data_fields'] = ['image', 'gt_bbox', 'gt_class', 'difficult']

    # 设置验证集相关的配置
    cfg_dict['EvalDataset'] = {}
    cfg_dict['EvalDataset']['name'] = 'VOCDataSet'
    cfg_dict['EvalDataset']['dataset_dir'] = user_cfg['EvalDataset']['dataset_dir']
    cfg_dict['EvalDataset']['anno_path'] = user_cfg['EvalDataset']['anno_path']
    cfg_dict['EvalDataset']['label_list'] = user_cfg['EvalDataset']['label_list']
    cfg_dict['EvalDataset']['data_fields'] = ['image', 'gt_bbox', 'gt_class', 'difficult']

    # 设置测试集相关的配置
    cfg_dict['TestDataset'] = {}
    cfg_dict['TestDataset']['name'] = 'ImageFolder'
    cfg_dict['TestDataset']['anno_path'] = os.path.join( user_cfg['TestDataset']['dataset_dir'], user_cfg['TestDataset']['label_list'] )

    # 修改评价指标相关配置
    cfg_dict['metric'] = 'VOC'
    # map_type (str): Calculation method of mean average
    # precision, currently support '11point' and
    # 'integral'. Default '11point'.
    cfg_dict['map_type'] = 'integral'
    cfg_dict['num_classes'] = num_classes

    print(json.dumps(cfg_dict))

    return num_classes, cfg_dict

def load_coco_dataset( user_cfg ):
    # 读取数据集的类别信息
    classes = user_cfg['classes']
    num_classes = len(classes)

    # 生成mmdetection对应的字典格式
    # 设置训练集相关的配置
    cfg_dict = {}
    cfg_dict['TrainDataset'] = {}
    cfg_dict['TrainDataset']['name'] = 'COCODataSet'
    cfg_dict['TrainDataset']['dataset_dir'] = user_cfg['TrainDataset']['dataset_dir']
    cfg_dict['TrainDataset']['image_dir'] = user_cfg['TrainDataset']['image_dir']
    cfg_dict['TrainDataset']['anno_path'] = user_cfg['TrainDataset']['anno_path']
    cfg_dict['TrainDataset']['data_fields'] = ['image', 'gt_bbox', 'gt_class', 'is_crowd']

    # 设置验证集相关的配置
    cfg_dict['EvalDataset'] = {}
    cfg_dict['EvalDataset']['name'] = 'COCODataSet'
    cfg_dict['EvalDataset']['dataset_dir'] = user_cfg['EvalDataset']['dataset_dir']
    cfg_dict['EvalDataset']['image_dir'] = user_cfg['EvalDataset']['image_dir']
    cfg_dict['EvalDataset']['anno_path'] = user_cfg['EvalDataset']['anno_path']
    cfg_dict['EvalDataset']['allow_empty'] = True

    # 设置测试集相关的配置
    cfg_dict['TestDataset'] = {}
    cfg_dict['TestDataset']['name'] = 'ImageFolder'
    cfg_dict['TestDataset']['dataset_dir'] = user_cfg['TestDataset']['dataset_dir']
    cfg_dict['TestDataset']['anno_path'] = user_cfg['TestDataset']['anno_path']

    # 修改评价指标相关配置
    # map_type (str): Calculation method of mean average
    # precision, currently support '11point' and
    # 'integral'. Default '11point'.
    # cfg_dict['map_type'] = 'integral'
    cfg_dict['metric'] = 'COCO'
    cfg_dict['num_classes'] = num_classes

    print(json.dumps(cfg_dict))

    return num_classes, cfg_dict
